package priv.pfz.raft.peer;

import com.google.common.collect.Lists;
import lombok.extern.slf4j.Slf4j;

import java.util.List;
import java.util.concurrent.*;
import java.util.function.Predicate;
import java.util.function.Supplier;

/**
 * 可自定义条件的CountDownLatch，监视多线程的执行情况，一旦满足条件或超时则返回。
 * 后续结果将被丢弃
 * @author pengfangzhou
 * @date 2022/2/27 2:02
 */
@Slf4j
public class ConditionLatch<T> {
    /**
     * 终止条件：对于当前已执行完成执行结果，判断是否终止
     * List<T>仅包含已执行完成的线程的结果
     * 符合终止条件前，每个线程完成时均会调用此函数
     */
    private final Predicate<List<T>> terminatePredicate;

    /**
     * 执行任务的线程池
     */
    private final Executor executor;

    private final List<T> resultList = Lists.newArrayList();
    private boolean terminated = false;
    private final BlockingQueue<List<T>> blockingQueue = new ArrayBlockingQueue<>(1);

    public ConditionLatch(Predicate<List<T>> terminatePredicate, Executor executor) {
        this.terminatePredicate = terminatePredicate;
        this.executor = executor;
    }

    /**
     * 添加一个任务
     * supplier必须自己处理异常
     */
    public void submit(Supplier<T> supplier) {
        CompletableFuture.supplyAsync(supplier, executor).thenAccept(this::testTerminate);
    }

    private synchronized void testTerminate(T result) {
        if (terminated) {
            log.info("discard result: {}", result);
            return;
        }
        resultList.add(result);
        if (terminatePredicate.test(resultList)) {
            terminated = true;
            log.info("offer results: {}", resultList);
            blockingQueue.offer(resultList);
        }
    }

    /**
     * 获取结果，超时抛出异常
     */
    public List<T> get(long timeoutMillis) {
        List<T> result;
        try {
            result = blockingQueue.poll(timeoutMillis, TimeUnit.MILLISECONDS);
        } catch (Exception e) {
            throw new RuntimeException("get exception", e);
        }
        if (result == null) {
            throw new RuntimeException("get timeout");
        }
        return result;
    }
}
